import asyncio
from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import List, Dict, Optional
from agent.tools.database_connection import DatabaseConnection  # 导入数据库连接管理类

# 定义SQL查询的输入模型
class SQLQuery(BaseModel):
    computation: str = Field(description="计算公式 (SQL 表达式)")
    data_source: str = Field(description="数据源表名")
    filters: Optional[List[Dict[str, str]]] = Field(default=None, description="过滤条件")
    group_by: Optional[List[str]] = Field(default=None, description="分组字段")
    
    class Config:
        schema_extra = {
            "example": {
                "computation": "SUM(approval_status) / COUNT(transaction_id)",
                "data_source": "transactions",
                "filters": [
                    {"field": "transaction_date", "condition": "BETWEEN '2025-06-01' AND '2025-06-30'"},
                    {"field": "region", "condition": "IN ('北京', '上海')"}
                ],
                "group_by": ["region", "industry"]
            }
        }

# SQL查询工具
@tool("sql_query", args_schema=SQLQuery, return_direct=True)
async def sql_query(computation: str, data_source: str, filters: Optional[List[Dict[str, str]]] = None, group_by: Optional[List[str]] = None) -> List[Dict]:
    """根据传入的计算公式、数据源、过滤条件生成SQL查询并返回计算结果"""
    try:
        # 使用数据库连接管理类获取数据库连接
        db_connection = DatabaseConnection.get_connection()
        if not db_connection:
            raise ValueError("Unable to connect to the database.")

        cursor = db_connection.cursor(dictionary=True)

        # 构建基本查询
        query = f"SELECT {computation} FROM {data_source}"

        # 构建 WHERE 子句（根据 filters 参数）
        where_clauses = []
        if filters:
            for filter in filters:
                field = filter.get("field")
                condition = filter.get("condition")
                if field and condition:
                    where_clauses.append(f"{field} {condition}")
            if where_clauses:
                query += " WHERE " + " AND ".join(where_clauses)

        # 构建 GROUP BY 子句（根据 group_by 参数）
        if group_by:
            query += f" GROUP BY {', '.join(group_by)}"
        
        # 执行查询并返回结果
        cursor.execute(query)
        result = cursor.fetchall()
        
        cursor.close()
        DatabaseConnection.close_connection(db_connection)
        
        return result

    except Exception as e:
        print(f"Error executing SQL query: {e}")
        return []
